# 加载数据集

import PIL.Image
import torch
from torch.utils.data import Dataset, DataLoader
import torchvision
import torchvision.transforms as transforms
import PIL
import numpy


def get_dataset(is_train=True, dowanload=True, dataset="cifar100", root="./data" ,is_transform="train"):

    # 进行图片变换
    if is_transform == "train":
        transform_dataset = transforms.Compose([
                                transforms.RandomCrop(32, padding=4),
                                transforms.RandomHorizontalFlip(),
                                transforms.RandomRotation(15),
                                transforms.ToTensor(),
                                # transforms.Normalize(mean, std)
                            ])
    else:
        transform_dataset = transforms.Compose([
                                transforms.ToTensor(),
                                # transforms.Normalize(mean, std)
                            ])
    
    assert dataset in ["cifar100", "cifar10"], "please choose the proper dataset!"     # 确保有合适的数据集 
    # 下载数据
    if dataset == "cifar100":
        res_dataset = torchvision.datasets.CIFAR100(root=root, train=is_train, 
                                                      download=dowanload, transform=transform_dataset)
    # image, label = res_dataset[12]
    # image.save("82.jpg")
    return res_dataset



def get_dataloader(dataset, shuffle=True, batch_size=512, numer_workers=16):
    res_loader = DataLoader(dataset, shuffle=shuffle, num_workers=numer_workers, 
                                      batch_size=batch_size)
    return dataset.classes, res_loader 


# 获取dataloader的均值和方差
def get_mean_std(mydataloader):

    # 累计各通道的和  与平方和
    total_sum = torch.zeros(3)
    total_squared_sum = torch.zeros(3)
    num_batches = 0
    
    for images, _ in mydataloader:
        # 将批次中每个图片shape 展开为 [batch_size, channels, height*width]
        images = images.view(images.shape[0], images.shape[1], -1)  # 参数 -1 自适应大小

        # 计算每个通道的均值和平方和
        total_sum += images.mean(dim=[0, 2])
        total_squared_sum += images.pow(2).mean(dim=[0, 2])
        num_batches += 1
    
    # 计算整个数据集的均值和标准差
    mean = total_sum / num_batches
    std = (total_squared_sum / num_batches - mean**2).sqrt()

    # 返回均值 和标准差
    return mean, std


# # 获取数据集
# def get_train_loader(mean, std, batch_size=1024, num_worker=16, shuffle=True):
    
#     # 进行数据预处理
#     transform_train = transforms.Compose([
#         transforms.RandomCrop(32, padding=4),
#         transforms.RandomHorizontalFlip(),
#         transforms.RandomRotation(15),
#         transforms.ToTensor(),
#         transforms.Normalize(mean, std)
#     ])

#     # 下载数据
#     cifar100_training = torchvision.datasets.CIFAR100(root="./data", train=True, 
#                                                       download=True, transform=transform_train)
#     cifar100_training_loader = DataLoader(cifar100_training, shuffle=shuffle, 
#                                           num_workers=num_worker, batch_size=batch_size)
    
#     # 返回数据集
#     return cifar100_training_loader


# # 获取测试集
# def get_val_loader(mean, std, batch_size=1024, numer_workers=16, shuffle=True):
#     transform_test = transforms.Compose([
#         transforms.ToTensor(),
#         transforms.Normalize(mean, std)
#     ])
#     cifar100_test = torchvision.datasets.CIFAR100(root="./data", train=False, 
#                                                   download=True, transform=transform_test)
#     cifar100_test_loader = DataLoader(cifar100_test, shuffle=shuffle, num_workers=numer_workers, 
#                                       batch_size=batch_size)
    
#     # 返回验证集
#     return cifar100_test_loader


if __name__ == "__main__":
    mean = torch.tensor([-0.2875, -0.2977, -0.3175])
    std = torch.tensor([0.3049, 0.2911, 0.2955])
    dataset = get_dataset()
    class_names, dataloader = get_dataloader(dataset)
    print(class_names[82])
    meean, std = get_mean_std(dataloader)
    print("均值和方差: ", mean, std)
    for image, label in dataloader:
        print(image.shape, label)
        break